package org.jiangwei.cmput696.entitylinking.algorithm;

import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import opennlp.tools.sentdetect.SentenceDetector;
import opennlp.tools.sentdetect.SentenceDetectorME;
import opennlp.tools.sentdetect.SentenceModel;
import opennlp.tools.tokenize.Tokenizer;
import opennlp.tools.tokenize.TokenizerME;
import opennlp.tools.tokenize.TokenizerModel;
import opennlp.tools.util.Span;

public class BagOfWords {

	private String wikiName;

	private String text;

	private String surface;

	private Map<String, Double> vector;

	private SentenceDetector sentenceDetector;

	private Tokenizer tokenizer;
	
	private List<String> seeds;

	private class TokenInfo {
		private String token;
		private int position;
		private int distToSurface = -1;
		private boolean isSurface;

		public String getToken() {
			return token;
		}

		public int getPosition() {
			return position;
		}

		public TokenInfo(String token, int position) {
			super();
			this.token = token;
			this.position = position;
		}

		public int getDistToSurface() {
			return distToSurface;
		}

		public void setDistToSurface(int distToSurface) {
			this.distToSurface = distToSurface;
		}

		public boolean isSurface() {
			return isSurface;
		}

		public void setSurface(boolean isSurface) {
			this.isSurface = isSurface;
		}

	}

	public BagOfWords(String wikiName, String text, String surface, List<String> seeds) {
		this.wikiName = wikiName;
		this.text = text;
		this.surface = surface;
		this.seeds = seeds;

		sentenceDetector = AlgoHelper.getSentenceDetector();
		tokenizer = AlgoHelper.getTokenizer();

		// tokenize text and build the vector (Map<String, int>)
		process(text, surface);
	}



	private void process(String text, String surface) {
		vector = new HashMap<String, Double>();

		// uses opennlp
		StringBuilder sb = new StringBuilder(text);
		Span[] sentSpans = sentenceDetector.sentPosDetect(text);

		int pos = 0;
		List<TokenInfo> tokenList = new ArrayList<TokenInfo>();
		
		Set<String> titleTokens = new HashSet<String>(0);
		boolean first = true;

		for (Span sentSpan : sentSpans) {
			String sent = sb.substring(sentSpan.getStart(), sentSpan.getEnd());
//			if (shouldInclude(surface, sent)) {
			if (shouldInclude(seeds, sent)) {
				for (String token : tokenizer.tokenize(sent)) {
					if (token.length() > 1 && token.charAt(0) >= 'A'
							&& token.charAt(0) <= 'Z') {
						++pos;
						TokenInfo ti = new TokenInfo(token, pos);
						
						token = token.toLowerCase();
						if (surface.toLowerCase().equals(token)) {
							ti.setSurface(true);
						}
						tokenList.add(ti);
						// if it is in the title
						if (first) {
							if (!titleTokens.contains(token)) {
								titleTokens.add(token);
							}
						}
					}
				}
				
				if (first) {
					first = false;
				}
			}
		}

		Map<String, Integer> countVector = new HashMap<String, Integer>();

		for (TokenInfo ti : tokenList) {
			String token = ti.getToken().toLowerCase();
			addToVector(countVector, token);
		}

		for (Map.Entry<String, Integer> countEntry : countVector.entrySet()) {
			double tf = countEntry.getValue();
			if (titleTokens.contains(countEntry.getKey())) {
				tf *= 1.5;
			}
			vector.put(countEntry.getKey(), tf);
		}
	}



	private boolean shouldInclude(String surface, String sent) {
		sent = sent.toLowerCase();
		String[] tokens = tokenizer.tokenize(surface);
		for (String token : tokens) {
			if (Character.isUpperCase(token.charAt(0)) && sent.contains(token.toLowerCase())) {
				return true;
			}
		}
		return false;
		
		/* TODO
		surface = surface.toLowerCase();
		sent = sent.toLowerCase();
		return sent.contains(surface) || surface.contains(sent);
		*/
	}
	
	private boolean shouldInclude(List<String> seeds, String sent) {
		sent = sent.toLowerCase();
		for (String token : seeds) {
			if (sent.contains(token)) {
				return true;
			}
		}
		return false;
	}

	public Map<String, Double> getVector() {
		return new HashMap<String, Double>(vector);
	}

	private void addToVector(Map<String, Integer> vector, String token) {
		token = token.toLowerCase();
		if (!vector.containsKey(token)) {
			vector.put(token, 1);
		} else {
			vector.put(token, vector.get(token) + 1);
		}
	}

	public String getWikiName() {
		return wikiName;
	}

	public String getText() {
		return text;
	}

	public String getSurface() {
		return surface;
	}
}
